import glob
import os
import json
import os.path
from tqdm import tqdm

import xml.sax
from shutil import copyfile
from xml import etree
import cv2
import numpy as np

from utils.plots import plot_one_box

try:
    import xml.etree.cElementTree as ET
except ImportError:
    import xml.etree.ElementTree as ET


def mkdir(path):  # 创建文件
    if not os.path.exists(path):
        os.makedirs(path)


def load_all_cvat_label(label_path):  # 加载cvat2d图像的标签
    cvat_label = dict()
    qianzhui_name_lists = []
    with open(label_path, "r") as f:
        tree = ET.parse(label_path)
        root = tree.getroot()
        image_node_list = root.findall("image")  # root中找到所有image关键词  所有节点找到
        bar = tqdm(image_node_list)
        for idx, image_node in enumerate(bar):
            bar.set_description("load_all_cvat_label: ----- {}/{}".format(idx + 1, len(bar)))
            image_dict = image_node.attrib
            image_name = image_dict["name"].split("/")[2].split(".")[0]  # 这个是图像的名称
            qianzhui_name_lists.append(image_name)

            box_node_list = image_node.findall("box")
            image_label_dict = dict()
            for box_node in box_node_list:

                box_dict = box_node.attrib

                try:
                    # image_label_dict['group_id'] = box_dict["group_id"]
                    image_label_dict[box_dict["group_id"]] = {"label": box_dict["label"], "xtl": box_dict["xtl"],
                                                              "ytl": box_dict["ytl"],
                                                              "xbr": box_dict["xbr"], "ybr": box_dict["ybr"]}  # 字典套字典
                except:
                    print(image_name + "  box no group id")

                # 内层字典key是label的名称  val还是一个字典  key是 坐标值索引  val是值

            cvat_label[image_name + 'box'] = image_label_dict  # 把标签的box信息字典放进字典里面去
            points_node_list = image_node.findall("points")
            points_label_dict = dict()
            for points_node in points_node_list:

                points_dict = points_node.attrib
                try:
                    # points_label_dict['group_id'] = points_dict['group_id']
                    points_label_dict[points_dict["group_id"]] = {"points": points_dict["points"]}  # 字典套字典
                except:
                    print(image_name + "   points no group id")

                # 内层字典key是label的名称  val还是一个字典  key是 坐标值索引  val是值

            cvat_label[image_name + 'points'] = points_label_dict  # 把标签的box信息字典放进字典里面去
        cvat_label['qianzhui'] = qianzhui_name_lists  # 信息名字前缀组成的list
        return cvat_label





# 合并cvat和lidar标签
def concat_lidar_cvat(images_labels_path, save_path):
    concat_label_path = save_path
    mkdir(concat_label_path)
    cvat_label_dict = load_all_cvat_label(images_labels_path)  # 加载并处理2d cvat的标签数据
    qianzhui_name_lists = cvat_label_dict['qianzhui']
    for qianzhui_name_list in qianzhui_name_lists:
        with open(os.path.join(save_path, qianzhui_name_list + '.txt'), 'a') as fp:
            for group_id, info in cvat_label_dict[qianzhui_name_list + 'box'].items():
                save_info_list = []
                cvat_n = info['label']
                save_info_list.append(cvat_n)
                xtl = float(info['xtl'])
                ytl = float(info['ytl'])
                xbr = float(info['xbr'])
                ybr = float(info['ybr'])
                x = str((xbr + xtl) / (1920 * 2))
                y = str((ybr + ytl) / (1080 * 2))
                w = str((xbr - xtl) / 1920)
                h = str((ybr - ytl) / 1080)
                if float(x) < 0 or float(y) < 0 or float(w) < 0 or float(h) < 0:
                    print(qianzhui_name_list + "xywh出错了！")
                save_info_list.append(x)
                save_info_list.append(y)
                save_info_list.append(w)
                save_info_list.append(h)
                try:
                    image_points_dict = cvat_label_dict[qianzhui_name_list + 'points'][group_id]
                    points = image_points_dict['points']
                    if cvat_n == '人' or cvat_n == '纸盒':
                        if len(points.split(',')) == 2:
                            yb = float(points.split(' ,')[1])
                        else:
                            print(qianzhui_name_list + "中的人 纸盒点标签有问题")
                    else:
                        if len(points.split(';')) == 2:
                            yb = (float(points.split(';')[0].split(',')[1]) + float(
                                points.split(';')[1].split(',')[1])) / 2
                        else:
                            print(qianzhui_name_list + "中的其他点标签有问题")
                except:
                    print(qianzhui_name_list + '    没有points与box对应')

                z = 3017.476 / abs(yb - 555.0000001) + 0.02548
                z = str(z)
                save_info_list.append(z)
                s = ' '.join(save_info_list)
                fp.write(s)
                fp.write('\n')




def start(images_labels_path, save_path):
    print("开始转换---------------------------->")
    concat_lidar_cvat(images_labels_path, save_path)
    print("转换完成！")


images_labels_path = "/run/user/1000/gvfs/sftp:host=10.10.10.114,user=fuyu/home/fuyu/data/cvat/YuanQu/2d_labels/annotations.xml"  # 图片的labels的地址

save_path = "/run/user/1000/gvfs/sftp:host=10.10.10.114,user=fuyu/home/fuyu/data/cvat/YuanQu/25d_labels"  # 最后保存得到的地址

start(images_labels_path=images_labels_path, save_path=save_path)
